"""Convert conversation-style JSONL data into Alpaca format entries."""

from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Iterable, List, Dict, Any


def parse_args() -> argparse.Namespace:
    """Configure and parse command line arguments."""
    parser = argparse.ArgumentParser(
        description="Convert a JSONL dataset into Alpaca system/instruction/output format."
    )
    parser.add_argument(
        "input_path",
        type=Path,
        help="Path to the source JSONL file with conversation data.",
    )
    parser.add_argument(
        "-o",
        "--output-path",
        type=Path,
        default=None,
        help="Destination JSONL path. Defaults to <input_name>_alpaca.jsonl next to the input file.",
    )
    return parser.parse_args()


def conversation_to_alpaca(system: str, conversation: Iterable[Dict[str, Any]]) -> List[Dict[str, str]]:
    """
    Expand a conversation into Alpaca-style records.

    Each assistant reply becomes one example paired with the most recent human instruction.
    """
    records: List[Dict[str, str]] = []
    last_human: str | None = None
    for turn in conversation:
        if "human" in turn:
            last_human = str(turn["human"])
        if "assistant" in turn:
            if last_human is None:
                # Skip assistant replies that do not have a preceding instruction.
                continue
            records.append(
                {
                    "system": system,
                    "instruction": last_human,
                    "output": str(turn["assistant"]),
                }
            )
    return records


def convert_file(input_path: Path, output_path: Path) -> None:
    """Read the source dataset and write Alpaca formatted JSONL."""
    alpaca_records: List[Dict[str, str]] = []
    with input_path.open("r", encoding="utf-8") as src:
        for line_number, line in enumerate(src, start=1):
            if not line.strip():
                continue
            try:
                payload = json.loads(line)
            except json.JSONDecodeError as error:
                raise ValueError(f"Line {line_number} is not valid JSON: {error}") from error
            system_prompt = str(payload.get("system", ""))
            conversation = payload.get("conversation", [])
            alpaca_records.extend(conversation_to_alpaca(system_prompt, conversation))

    with output_path.open("w", encoding="utf-8") as dst:
        for record in alpaca_records:
            dst.write(json.dumps(record, ensure_ascii=False) + "\n")


def main() -> None:
    """Entrypoint for the command line tool."""
    args = parse_args()
    input_path: Path = args.input_path
    output_path: Path = (
        args.output_path
        if args.output_path is not None
        else input_path.with_name(f"{input_path.stem}_alpaca.jsonl")
    )
    convert_file(input_path, output_path)


if __name__ == "__main__":
    main()
